import numpy as np
import matplotlib.pyplot as plt
import pywt
from scipy.signal import stft, hilbert
from scipy.fftpack import fft

# 生成模拟数据
fs = 1000  # 采样频率
t = np.linspace(0, 1, fs, endpoint=False)
amplitude = 1.5  # 基础振幅
np.random.seed(0)  # 确保每次运行代码生成相同的随机数

# 创建更复杂的信号
signal = sum(amplitude * np.random.rand() * np.sin(2 * np.pi * np.random.randint(30, 100) * t + np.random.rand())
             for _ in range(5))  # 随机生成5个正弦波的和

noise = np.random.normal(0, 0.2, fs)  # 高斯噪声

# 生成白噪声
noise_amplitude = 0.5  # 噪声的振幅
#noise = noise_amplitude * np.random.normal(size=t.shape)  # 均值为0，标准差为0.5的正态分布
mixed_signal = signal + noise  # 含噪声的信号

# 定义信噪比计算函数
def calculate_snr(signal, noise):
    signal_power = np.mean(signal**2)
    noise_power = np.mean(noise**2)
    return 10 * np.log10(signal_power / noise_power)

# 无干扰情况下的处理
snr_no_interference = calculate_snr(signal, noise)

# 加白噪声情况下的处理
snr_white_noise = calculate_snr(mixed_signal, noise)

# 相似性检测法的处理
# TODO: 添加相似性检测法的信号处理代码和计算信噪比的代码
# 第一阶段处理：高通滤波和小波变换去噪

# # 高通滤波
# highpass_signal = mixed_signal - np.mean(mixed_signal)
# # 小波变换去噪
# coeffs = pywt.wavedec(highpass_signal, 'db4', level=4)
# threshold = 0.1
# thresholded_coeffs = [pywt.threshold(i, threshold, mode='soft') for i in coeffs]
# wavelet_signal = pywt.waverec(thresholded_coeffs, 'db4')
#
# # 第二阶段处理：FFT和希尔伯特变换
# # FFT
# fft_signal = fft(wavelet_signal)
# # SFFT和希尔伯特变换
# f, t_, Zxx = stft(wavelet_signal, fs, nperseg=100)
# hilbert_signal = hilbert(wavelet_signal)
# envelope_signal = np.abs(hilbert_signal)




# 预测性维护系统的处理
# TODO: 添加预测性维护系统的信号处理代码和计算信噪比的代码
#第一阶段处理：高通滤波、低通滤波和小波变换去噪
#高通滤波
highpass_signal = mixed_signal - np.mean(mixed_signal)
#低通滤波
lowpass_signal = np.convolve(highpass_signal, np.ones(100)/100, mode='same')
#小波变换去噪
coeffs = pywt.wavedec(lowpass_signal, 'db4', level=4)
threshold = 0.1
thresholded_coeffs = [pywt.threshold(i, threshold, mode='soft') for i in coeffs]
wavelet_signal = pywt.waverec(thresholded_coeffs, 'db4')

#第二阶段处理：FFT和希尔伯特变换
#FFT
fft_signal = fft(wavelet_signal)
#SFFT和希尔伯特变换
f, t_, Zxx = stft(wavelet_signal, fs, nperseg=100)
hilbert_signal = hilbert(wavelet_signal)
envelope_signal = np.abs(hilbert_signal)




#信号功率
signal_power = np.mean(signal**2)
#噪声功率
#noise_power = np.mean(noise**2)
#白噪声功率
noise_power = np.mean(noise**2)

#信噪比
print(f"Signal Power: {signal_power:.2f}")
print(f"Noise Power: {noise_power:.2f}")



snr_wavelet = calculate_snr(wavelet_signal, mixed_signal - wavelet_signal)  # 计算信噪比
snr_envelope = calculate_snr(envelope_signal, mixed_signal - envelope_signal)  # 计算信噪比

#

# 输出相似性检测法和预测性维护系统的信噪比
print(f"SNR of original signal: {snr_no_interference:.2f} dB")
print(f"SNR after adding white noise: {snr_white_noise:.2f} dB")
print(f"SNR after similarity detection: {snr_wavelet:.2f} dB")
print(f"SNR after predictive maintenance system: {snr_envelope:.2f} dB")

